from PyQt5.QtCore import QThread
import pathlib
from PIL import Image
from xml.dom.minidom import parse, Text
import shutil
import itertools
import yaml


def xml2dict(xml):
    def dfs_parse(dom, data):
        for child in dom.childNodes:
            if len(child.childNodes) == 1:  # 文本节点
                data[child.tagName] = child.firstChild.nodeValue
            elif not isinstance(child, Text):
                ans = {}
                if child.tagName not in data:
                    data[child.tagName] = ans
                elif isinstance(data[child.tagName], dict):
                    data[child.tagName] = [data[child.tagName], ans]
                else:
                    data[child.tagName].append(ans)
                dfs_parse(child, ans)
        return data

    return dfs_parse(parse(xml), {})


def parse_xml(xml_path, class_to_idx):
    xml_path = str(xml_path)
    warn = []
    boxes, labels = [], []
    objects = xml2dict(xml_path)["annotation"]
    if "object" not in objects:
        warn.append(f"{xml_path} 不存在物体")
        return boxes, labels, warn
    width = int(objects["size"]["width"])
    height = int(objects["size"]["height"])
    objects = objects["object"]
    objects = objects if isinstance(objects, list) else [objects]
    for obj in objects:
        label_name = obj["name"]
        #################################
        # 过滤条件
        # if (label_name == "face" or 
        #     width * height < 250*250):
        #     continue
        #################################
        if label_name in class_to_idx:
            label_id = class_to_idx[label_name]
        else:
            label_id = len(class_to_idx)
            class_to_idx[label_name] = label_id
        x1 = int(obj["bndbox"]["xmin"])
        y1 = int(obj["bndbox"]["ymin"])
        x2 = int(obj["bndbox"]["xmax"])
        y2 = int(obj["bndbox"]["ymax"])
        labels.append(label_id)
        boxes.append([x1, y1, x2, y2])
    return boxes, labels, warn


def convert(img_iter, train_rate, output_path, class_to_idx, log, set_progress):
    warn = []
    output_path = pathlib.Path(output_path)
    r1 = output_path / "train" / "images"
    r2 = output_path / "train" / "labels"
    r3 = output_path / "val" / "images"
    r4 = output_path / "val" / "labels"
    for sub_dir in [r1, r2, r3, r4]:
        if not sub_dir.exists():
            sub_dir.mkdir(parents=True)

    idx, val_idx = 0, int(1 / (1 - train_rate))
    img_list = list(img_iter)
    total = len(img_list)
    for n, img_path in enumerate(img_list):
        xml_path = img_path.parent.parent / "Annotations" / (img_path.stem + ".xml")
        if not xml_path.exists():
            warn.append(f"{xml_path} 不存在")
        else:
            boxes, labels, w = parse_xml(xml_path, class_to_idx)
            warn.extend(w)
            if len(boxes) == 0:
                continue

            sub_dir_name = "train" if idx % val_idx != 0 else "val"
            idx += 1
            name = str(idx)

            o = output_path / sub_dir_name / "images" / (name + ".jpg")
            # copy image
            shutil.copy(img_path, o)
            log(f"复制 {img_path} 到 {o}")
            set_progress(int((n + 1) / total * 100))
            # create label
            img = Image.open(img_path)
            with open((output_path / sub_dir_name / "labels" / (str(idx) + ".txt")), "w") as f:
                for b, l in zip(boxes, labels):
                    x1, y1, x2, y2 = b
                    cx, cy, w, h = (x1 + x2) / 2, (y1 + y2) / \
                                   2, x2 - x1, y2 - y1
                    cx, cy, w, h = cx / img.width, cy / img.height, w / img.width, h / img.height
                    f.write(f"{l} {cx} {cy} {w} {h}\n")
    return warn


class ConvertThread(QThread):
    def __init__(self, parent, img_dir_list, train_rate, output_path):
        super().__init__(parent=parent)
        self.parent = parent
        self.log = parent.output.emit
        self.set_process = parent.progress.emit
        self.img_dir_list = img_dir_list
        self.train_rate = train_rate
        self.output_path = pathlib.Path(output_path)

    def run(self):
        iterlist = []
        for img_dir in self.img_dir_list:
            # 获取所有图片
            iterlist.append(pathlib.Path(img_dir).glob("*.jpg"))
        img_iter = itertools.chain(*iterlist)
        class_to_idx = {}
        warn = convert(img_iter, self.train_rate, self.output_path, class_to_idx, self.log, self.set_process)
        idx_to_class = {v: k for k, v in class_to_idx.items()}
        data_yaml = {
            "path": str(self.output_path),
            "names": idx_to_class,
            "train": "train/images",
            "val": "val/images",
        }
        dir_name = self.output_path.name
        kaggle_data_yaml = {
            "path": f"/kaggle/input/{dir_name.replace('_', '-')}",
            "names": idx_to_class,
            "train": "train/images",
            "val": "val/images",
        }
        with open(self.output_path / "data.yaml", "w") as f:
            yaml.dump(data_yaml, f)
        with open(self.output_path / "kaggle_data.yaml", "w") as f:
            yaml.dump(kaggle_data_yaml, f)

        self.log("======================")
        self.log("转换完成")
        for w in warn:
            self.log(w)

        cmd1 = "!git clone https://github.com/ultralytics/yolov5.git && wandb off"
        cmd2 = "!pip install ultralytics"
        cmd3 = f"!python yolov5/train.py --img 640 --epochs 10 --data /kaggle/input/{dir_name}/kaggle_data.yaml --weights yolov5s.pt"
        with open(self.output_path / "cmd.txt", "w") as f:
            f.write(cmd1 + "\n")
            f.write(cmd2 + "\n")
            f.write(cmd3 + "\n")
        self.parent.enable_ui()
